from .base_exchange import *
import okx.PublicData as PublicAPI
import okx.Account as Account
import okx.Trade as Trade
from okx.websocket import WebSocketAPI
import asyncio
import time

class OKXExchange(BaseExchange):
    def __init__(self, api_key: str = "", api_secret: str = "", passphrase: str = "", 
                 is_futures: bool = False, testnet: bool = False):
        super().__init__(api_key, api_secret, passphrase)
        self.name = "OKX"
        self.is_futures = is_futures
        self.testnet = testnet
        
        # API配置
        self.api_config = {
            "api_key": api_key,
            "api_secret_key": api_secret,
            "passphrase": passphrase,
            "flag": "0" if not testnet else "1",  # 0: 生产环境, 1: 模拟环境
            "debug": False
        }
        
        # 初始化REST API客户端
        self.public_api = PublicAPI.PublicAPI(**self.api_config)
        self.account_api = Account.AccountAPI(**self.api_config)
        self.trade_api = Trade.TradeAPI(**self.api_config)
        
        # WebSocket连接
        self.ws = None
        self.ws_connected = False
        
        # 缓存数据
        self.tickers = {}
        self.orderbooks = {}
        self.balances = {}
        self.positions = {}
        
    async def connect_ws(self):
        """连接WebSocket"""
        try:
            # 初始化WebSocket
            self.ws = WebSocketAPI(
                api_key=self.api_key,
                api_secret_key=self.api_secret,
                passphrase=self.passphrase,
                test=self.testnet
            )
            
            # 设置回调函数
            self.ws.start(
                channel_type=['public', 'private'],
                callback=self._handle_ws_message
            )
            
            # 订阅频道
            await self._subscribe_channels()
            
            self.ws_connected = True
            
        except Exception as e:
            raise Exception(f"Failed to connect WebSocket: {str(e)}")
            
    async def _subscribe_channels(self):
        """订阅WebSocket频道"""
        try:
            # 订阅行情数据
            self.ws.subscribe([
                {
                    "channel": "tickers",
                    "instId": "BTC-USDT" if not self.is_futures else "BTC-USDT-SWAP"
                },
                {
                    "channel": "books",
                    "instId": "BTC-USDT" if not self.is_futures else "BTC-USDT-SWAP"
                }
            ])
            
            # 订阅账户和持仓数据
            if self.api_key:
                self.ws.subscribe([
                    {
                        "channel": "account",
                        "instType": "SWAP" if self.is_futures else "SPOT"
                    },
                    {
                        "channel": "positions",
                        "instType": "SWAP" if self.is_futures else "SPOT"
                    }
                ])
                
        except Exception as e:
            raise Exception(f"Failed to subscribe channels: {str(e)}")
            
    def _handle_ws_message(self, message):
        """处理WebSocket消息"""
        try:
            if 'event' in message:
                return  # 忽略事件消息
                
            channel = message.get('arg', {}).get('channel')
            if not channel:
                return
                
            if channel == 'tickers':
                self._handle_ticker(message)
            elif channel == 'books':
                self._handle_orderbook(message)
            elif channel == 'account':
                self._handle_account(message)
            elif channel == 'positions':
                self._handle_position(message)
                
        except Exception as e:
            print(f"Error handling message: {e}")
            
    def _handle_ticker(self, message):
        """处理行情数据"""
        try:
            data = message['data'][0]
            symbol = data['instId']
            self.tickers[symbol] = {
                'last': float(data['last']),
                'bid': float(data['bidPx']),
                'ask': float(data['askPx']),
                'volume': float(data['vol24h']),
                'timestamp': int(data['ts'])
            }
        except Exception as e:
            print(f"Error handling ticker: {e}")
            
    def _handle_orderbook(self, message):
        """处理深度数据"""
        try:
            data = message['data'][0]
            symbol = data['instId']
            self.orderbooks[symbol] = {
                'asks': [(float(price), float(qty)) for price, qty in data['asks']],
                'bids': [(float(price), float(qty)) for price, qty in data['bids']],
                'timestamp': int(data['ts'])
            }
        except Exception as e:
            print(f"Error handling orderbook: {e}")
            
    async def get_ticker(self, symbol: str) -> Ticker:
        """获取最新行情"""
        try:
            response = self.public_api.get_ticker(
                instId=symbol
            )
            
            if response['code'] != '0':
                raise Exception(response['msg'])
                
            ticker_data = response['data'][0]
            return Ticker(
                symbol=symbol,
                last_price=float(ticker_data['last']),
                bid_price=float(ticker_data['bidPx']),
                ask_price=float(ticker_data['askPx']),
                volume_24h=float(ticker_data['vol24h']),
                timestamp=int(ticker_data['ts'])
            )
        except Exception as e:
            raise Exception(f"Failed to get ticker: {str(e)}")
            
    async def get_orderbook(self, symbol: str, limit: int = 20) -> OrderBook:
        """获取订单簿"""
        try:
            response = self.public_api.get_orderbook(
                instId=symbol,
                sz=str(limit)
            )
            
            if response['code'] != '0':
                raise Exception(response['msg'])
                
            data = response['data'][0]
            return OrderBook(
                asks=[(float(price), float(qty)) for price, qty in data['asks']],
                bids=[(float(price), float(qty)) for price, qty in data['bids']],
                timestamp=int(data['ts'])
            )
        except Exception as e:
            raise Exception(f"Failed to get orderbook: {str(e)}")
            
    async def get_balance(self, asset: str = "") -> Dict[str, Balance]:
        """获取账户余额"""
        try:
            response = self.account_api.get_account_balance()
            
            if response['code'] != '0':
                raise Exception(response['msg'])
                
            balances = {}
            for detail in response['data'][0]['details']:
                ccy = detail['ccy']
                if not asset or ccy == asset:
                    balances[ccy] = Balance(
                        asset=ccy,
                        free=float(detail['availBal']),
                        locked=float(detail['frozenBal']),
                        total=float(detail['cashBal']),
                        timestamp=time.time()
                    )
            return balances
        except Exception as e:
            raise Exception(f"Failed to get balance: {str(e)}")
            
    async def get_positions(self, symbol: str = "") -> Dict[str, Position]:
        """获取持仓信息"""
        try:
            if not self.is_futures:
                return {}
                
            response = self.account_api.get_positions(
                instType="SWAP",
                instId=symbol if symbol else None
            )
            
            if response['code'] != '0':
                raise Exception(response['msg'])
                
            positions = {}
            for pos in response['data']:
                symbol = pos['instId']
                positions[symbol] = Position(
                    symbol=symbol,
                    size=float(pos['pos']),
                    entry_price=float(pos['avgPx']),
                    leverage=float(pos['lever']),
                    margin=float(pos['margin']),
                    unrealized_pnl=float(pos['upl']),
                    liquidation_price=float(pos['liqPx']),
                    timestamp=time.time()
                )
            return positions
        except Exception as e:
            raise Exception(f"Failed to get positions: {str(e)}")
            
    async def create_market_order(self, symbol: str, side: OrderSide, amount: float) -> Order:
        """创建市价单"""
        try:
            if self.is_futures:
                response = self.trade_api.place_order(
                    instId=symbol,
                    tdMode='cross',  # 全仓模式
                    side=side.value,
                    ordType='market',
                    sz=str(amount),
                    posSide='long' if side == OrderSide.BUY else 'short'
                )
            else:
                response = self.trade_api.place_order(
                    instId=symbol,
                    tdMode='cash',  # 现货
                    side=side.value,
                    ordType='market',
                    sz=str(amount)
                )
                
            if response['code'] != '0':
                raise Exception(response['msg'])
                
            order_id = response['data'][0]['ordId']
            return await self.get_order(symbol, order_id)
            
        except Exception as e:
            raise Exception(f"Failed to create market order: {str(e)}")
            
    async def create_limit_order(self, symbol: str, side: OrderSide, 
                               price: float, amount: float) -> Order:
        """创建限价单"""
        try:
            if self.is_futures:
                response = self.trade_api.place_order(
                    instId=symbol,
                    tdMode='cross',
                    side=side.value,
                    ordType='limit',
                    px=str(price),
                    sz=str(amount),
                    posSide='long' if side == OrderSide.BUY else 'short'
                )
            else:
                response = self.trade_api.place_order(
                    instId=symbol,
                    tdMode='cash',
                    side=side.value,
                    ordType='limit',
                    px=str(price),
                    sz=str(amount)
                )
                
            if response['code'] != '0':
                raise Exception(response['msg'])
                
            order_id = response['data'][0]['ordId']
            return await self.get_order(symbol, order_id)
            
        except Exception as e:
            raise Exception(f"Failed to create limit order: {str(e)}")
            
    async def set_leverage(self, symbol: str, leverage: int) -> bool:
        """设置杠杆倍数"""
        try:
            if not self.is_futures:
                return False
                
            response = self.account_api.set_leverage(
                instId=symbol,
                lever=str(leverage),
                mgnMode='cross'  # 全仓模式
            )
            
            return response['code'] == '0'
            
        except Exception as e:
            raise Exception(f"Failed to set leverage: {str(e)}")
            
    async def get_funding_rate(self, symbol: str) -> float:
        """获取资金费率"""
        try:
            if not self.is_futures:
                return 0.0
                
            response = self.public_api.get_funding_rate(instId=symbol)
            
            if response['code'] != '0':
                raise Exception(response['msg'])
                
            return float(response['data'][0]['fundingRate'])
            
        except Exception as e:
            raise Exception(f"Failed to get funding rate: {str(e)}")
            
    async def get_funding_time(self, symbol: str) -> int:
        """获取下次资金费时间"""
        try:
            if not self.is_futures:
                return 0
                
            response = self.public_api.get_funding_time(instId=symbol)
            
            if response['code'] != '0':
                raise Exception(response['msg'])
                
            return int(response['data'][0]['fundingTime'])
            
        except Exception as e:
            raise Exception(f"Failed to get funding time: {str(e)}")
            
    def _convert_status(self, status: str) -> OrderStatus:
        """转换订单状态"""
        status_map = {
            'live': OrderStatus.OPEN,
            'filled': OrderStatus.FILLED,
            'canceled': OrderStatus.CANCELED,
            'rejected': OrderStatus.REJECTED
        }
        return status_map.get(status.lower(), OrderStatus.PENDING)
        
    async def _handle_order_update(self, message: dict):
        """处理订单更新消息"""
        try:
            data = message['data'][0]
            order = Order(
                exchange_id=self.name,
                order_id=data['ordId'],
                symbol=data['instId'],
                order_type=OrderType.MARKET if data['ordType'] == 'market' else OrderType.LIMIT,
                side=OrderSide.BUY if data['side'] == 'buy' else OrderSide.SELL,
                price=float(data['px']) if data['px'] else 0,
                amount=float(data['sz']),
                filled=float(data['accFillSz']),
                status=self._convert_status(data['state']),
                timestamp=int(data['uTime'])
            )
            
            # 更新订单缓存
            if order.status in [OrderStatus.FILLED, OrderStatus.CANCELED, OrderStatus.REJECTED]:
                self.orders.pop(order.order_id, None)
            else:
                self.orders[order.order_id] = order
                
        except Exception as e:
            self.logger.error(f"Error handling order update: {str(e)}")
            
    async def cancel_order(self, symbol: str, order_id: str) -> bool:
        """取消订单"""
        try:
            response = self.trade_api.cancel_order(
                instId=symbol,
                ordId=order_id
            )
            
            if response['code'] != '0':
                raise Exception(response['msg'])
                
            return True
            
        except Exception as e:
            raise Exception(f"Failed to cancel order: {str(e)}")
            
    async def get_order(self, symbol: str, order_id: str) -> Order:
        """获取订单信息"""
        try:
            response = self.trade_api.get_order(
                instId=symbol,
                ordId=order_id
            )
            
            if response['code'] != '0':
                raise Exception(response['msg'])
                
            data = response['data'][0]
            return Order(
                exchange_id=self.name,
                order_id=order_id,
                symbol=symbol,
                order_type=OrderType.MARKET if data['ordType'] == 'market' else OrderType.LIMIT,
                side=OrderSide.BUY if data['side'] == 'buy' else OrderSide.SELL,
                price=float(data['px']) if data['px'] else 0,
                amount=float(data['sz']),
                filled=float(data['accFillSz']),
                status=self._convert_status(data['state']),
                timestamp=int(data['uTime'])
            )
            
        except Exception as e:
            raise Exception(f"Failed to get order: {str(e)}")
            
    async def get_open_orders(self, symbol: str = "") -> List[Order]:
        """获取未完成订单"""
        try:
            response = self.trade_api.get_order_list(
                instType="SWAP" if self.is_futures else "SPOT",
                instId=symbol if symbol else None,
                state="live"
            )
            
            if response['code'] != '0':
                raise Exception(response['msg'])
                
            orders = []
            for data in response['data']:
                orders.append(Order(
                    exchange_id=self.name,
                    order_id=data['ordId'],
                    symbol=data['instId'],
                    order_type=OrderType.MARKET if data['ordType'] == 'market' else OrderType.LIMIT,
                    side=OrderSide.BUY if data['side'] == 'buy' else OrderSide.SELL,
                    price=float(data['px']) if data['px'] else 0,
                    amount=float(data['sz']),
                    filled=float(data['accFillSz']),
                    status=self._convert_status(data['state']),
                    timestamp=int(data['uTime'])
                ))
            return orders
            
        except Exception as e:
            raise Exception(f"Failed to get open orders: {str(e)}")
            
    async def _handle_account(self, message: dict):
        """处理账户更新消息"""
        try:
            data = message['data'][0]
            for balance in data['details']:
                ccy = balance['ccy']
                self.balances[ccy] = Balance(
                    asset=ccy,
                    free=float(balance['availBal']),
                    locked=float(balance['frozenBal']),
                    total=float(balance['cashBal']),
                    timestamp=time.time()
                )
        except Exception as e:
            self.logger.error(f"Error handling account update: {str(e)}")
            
    async def _handle_position(self, message: dict):
        """处理持仓更新消息"""
        try:
            data = message['data'][0]
            symbol = data['instId']
            self.positions[symbol] = Position(
                symbol=symbol,
                size=float(data['pos']),
                entry_price=float(data['avgPx']),
                leverage=float(data['lever']),
                margin=float(data['margin']),
                unrealized_pnl=float(data['upl']),
                liquidation_price=float(data['liqPx']),
                timestamp=time.time()
            )
        except Exception as e:
            self.logger.error(f"Error handling position update: {str(e)}")
            
    async def _handle_orderbook(self, message: dict):
        """处理深度数据更新"""
        try:
            data = message['data'][0]
            symbol = data['instId']
            self.orderbooks[symbol] = OrderBook(
                asks=[(float(price), float(qty)) for price, qty in data['asks']],
                bids=[(float(price), float(qty)) for price, qty in data['bids']],
                timestamp=int(data['ts'])
            )
        except Exception as e:
            self.logger.error(f"Error handling orderbook update: {str(e)}") 